#include <iostream>
#include <fstream>
#include <sstream>
#include <map>
#include <vector>
#include "term.hpp"
#include "doc.hpp"

int main(int argc, char* argv[]){
	//counters
	int p=0,c=0,f=0,s=0;
	std::string line;
	std::stringstream ss;
	std::ifstream training;
	std::ifstream test;
	std::ofstream output;
	std::map<std::string,Term> terms;
	std::vector<Doc> docs;
	if(argc!=3){
		std::cout << "Wrong # of Args!" << std::endl;
		return 1;
	}
	training.open(argv[1]);
	test.open(argv[2]);
	if (training.is_open() && test.is_open()){
		//parse training set and count terms
		while ( training.good() ){
			ss.str("");
			ss.clear();
			std::getline (training,line);
			ss << line;
			ss >> line;
			if(line=="project"){
				p++;
				while(ss >> line){
					terms[line].p++;
				}
			}
			else if(line=="course"){
				c++;
				while(ss >> line){
					terms[line].c++;
				}
			}
			else if(line=="faculty"){
				f++;
				while(ss >> line){
					terms[line].f++;
				}
			}
			else if(line=="student"){
				s++;
				while(ss >> line){
					terms[line].s++;
				}
			}
		}
		//update terms
		for(std::map<std::string,Term>::iterator i=terms.begin();i!=terms.end();i++){
			(i->second).update(p+terms.size(),c+terms.size(),f+terms.size(),s+terms.size());
			/*
			std::cout << "term = " << i->first << std::endl;
			std::cout << "pp = " << (i->second).ppt << std::endl;
			std::cout << "pc = " << (i->second).pct << std::endl;
			std::cout << "pf = " << (i->second).pft << std::endl;
			std::cout << "ps = " << (i->second).pst << std::endl;
			*/
		}
		//calculate probabilities
		Term::pp = ((double)p)/(p+c+f+s);
		Term::pc = ((double)c)/(p+c+f+s);
		Term::pf = ((double)f)/(p+c+f+s);
		Term::ps = ((double)s)/(p+c+f+s);
		training.close();
		//start classifying
		while(test.good()){
			ss.str("");
			ss.clear();
			std::getline(test,line);
			ss << line;
			docs.push_back(Doc());
			while(ss >> line){
				terms[line].update(p+terms.size(),c+terms.size(),f+terms.size(),s+terms.size());
				docs.back().addTerm(terms[line]);
			}
			for(std::map<std::string,Term>::iterator i=terms.begin();i!=terms.end();i++){
				(i->second).added=false;
			}
		}
		test.close();
		//get rid of extra blank doc
		docs.pop_back();
		//print results
		output.open("predictions.txt");
		if(output.is_open()){
			for(std::vector<Doc>::iterator i=docs.begin();i!=docs.end();i++){
				i->print(output);
			}
		}
		output.close();
	}
	else{
		std::cout << "Unable to open files" << std::endl;
		return 2;
	}
	//std::system("pause");
	return 0;
}